import os
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.nn import SAGPooling, GCNConv, GATConv, global_mean_pool, global_max_pool
from model import *

class STTraceAD(nn.Module):

    def __init__(self, args, embed_model=None):
        super(STTraceAD, self).__init__()

        self.window = args.window
        self.span_feats = args.span_feats
        self.metric_feats = args.metric_feats

        self.base_dir = args.base_dir
        self.model_fname = args.model_fname

        self.embed_model = embed_model
        self.temporal_model = TemporalModel(args.span_feats, args.rnn_hidden, args.rnn_layer, args.rnn)
        self.aggr_model = AggrModel(args.env_feats, args.embed_size, args.aggr_hidden, args.aggr_layer)
        self.stgnn = STGNN(args.env_feats+args.te_feats, args.gnn_hidden, args.pooling_ratio)
        self.detector = Detector(args.st_feats*2, args.dropout)

    def forward(self, data):
        # in.x.shape: (~B*L, F_s + w*F_m)
        # out.shape: (B, 2)

        x, edge_index, batch = data.x, data.edge_index, data.batch
        span, metric = x[:, :self.span_feats], x[:, self.span_feats:]

        # metric embedding and aggregation
        metric = metric.reshape(-1, self.window, self.metric_feats)
        if self.embed_model is not None:
            metric_embeded = self.embed_model(metric)  # (~BL, w, H_e)
        else:
            metric_embeded = metric
        x_env = self.aggr_model(metric_embeded)        # (~BL, H_en)

        # temporal modeling
        x_te = self.temporal_model(span, batch)       # (~BL, H_te)

        # spatial-temporal modeling
        x = torch.concat([x_te, x_env], dim=1)        # (~BL, H_te + H_en)
        x_st = self.stgnn(x, edge_index, batch)       # (B, 2*H_st)

        # anomaly detection
        y = self.detector(x_st)                     # (B, 2)

        return -y

    def load(self):
        file_path = os.path.join(self.base_dir, self.model_fname)
        self.load_state_dict(torch.load(file_path))

    def save(self):
        file_path = os.path.join(self.base_dir, self.model_fname)
        if os.path.exists(file_path):
            os.remove(file_path)
        torch.save(self.state_dict(), file_path)


class STTraceADV2(nn.Module):
    '''
    temporal spatial, no metric embedding and aggregation
    '''
    def __init__(self, args):
        super(STTraceADV2, self).__init__()

        self.base_dir = args.base_dir
        self.model_fname = args.model_fname

        self.temporal_model = TemporalModel(args.span_feats, args.rnn_hidden, args.rnn_layer, args.rnn)
        self.stgnn = STGNN(args.te_feats, args.gnn_hidden, args.pooling_ratio)
        self.detector = Detector(args.st_feats*2, args.dropout)

    def forward(self, data):
        # in.x.shape: (~B*L, F_s)
        # out.shape: (B, 2)

        x, edge_index, batch = data.x, data.edge_index, data.batch

        # temporal modeling
        x_te = self.temporal_model(x, batch)       # (~BL, H_te)

        # spatial-temporal modeling
        x_st = self.stgnn(x_te, edge_index, batch)       # (B, 2*H_st)

        # anomaly detection
        y = self.detector(x_st)                     # (B, 2)

        return y

    def load(self):
        file_path = os.path.join(self.base_dir, self.model_fname)
        self.load_state_dict(torch.load(file_path))

    def save(self):
        file_path = os.path.join(self.base_dir, self.model_fname)
        if os.path.exists(file_path):
            os.remove(file_path)
        torch.save(self.state_dict(), file_path)


class STTraceADV3(nn.Module):
    '''
    spatial only, no temporal and metric embedding
    '''
    def __init__(self, args):
        super(STTraceADV3, self).__init__()

        self.base_dir = args.base_dir
        self.model_fname = args.model_fname

        self.stgnn = STGNN(args.span_feats, args.gnn_hidden, args.pooling_ratio)
        self.detector = Detector(args.st_feats*2, args.dropout)

    def forward(self, data):
        # in.x.shape: (~B*L, F_s)
        # out.shape: (B, 2)

        x, edge_index, batch = data.x, data.edge_index, data.batch

        # spatial modeling
        x = self.stgnn(x, edge_index, batch)       # (B, 2*H_st)

        # anomaly detection
        y = self.detector(x)                     # (B, 2)

        return y

    def load(self):
        file_path = os.path.join(self.base_dir, self.model_fname)
        self.load_state_dict(torch.load(file_path))

    def save(self):
        file_path = os.path.join(self.base_dir, self.model_fname)
        if os.path.exists(file_path):
            os.remove(file_path)
        torch.save(self.state_dict(), file_path)


class GATAD(nn.Module):
    def __init__(self, args):
        super(GATAD, self).__init__()

        self.base_dir = args.base_dir
        self.model_fname = args.model_fname

        self.embed = nn.Linear(args.span_feats, 32)
        self.gnn1 = GATConv(32, 64, 1)
        self.gnn2 = GATConv(64, 64, 1)
        self.gnn3 = GATConv(64, 32, 1)
        self.detector = Detector(32, args.dropout)

    def forward(self, data):
        # in.x.shape: (~B*L, F_s)
        # out.shape: (B, 2)

        x, edge_index, batch = data.x, data.edge_index, data.batch

        x = self.embed(x)
        x = self.gnn1(x, edge_index)
        x = self.gnn2(x, edge_index)
        x = self.gnn3(x, edge_index)
        x = global_max_pool(x, batch)

        # anomaly detection
        y = self.detector(x)                     # (B, 2)

        return y

    def load(self):
        file_path = os.path.join(self.base_dir, self.model_fname)
        self.load_state_dict(torch.load(file_path))

    def save(self):
        file_path = os.path.join(self.base_dir, self.model_fname)
        if os.path.exists(file_path):
            os.remove(file_path)
        torch.save(self.state_dict(), file_path)